import numpy as np
import matplotlib.pylab as mp
from math import exp
## CI dans un tableau ##


## METHODE D'EULER

def step_euler(y,t,h,F):
    n=np.shape(y)
    z=np.zeros(n)
    z=y+h*F(y,t)
    return z

## METHODE  DU POINT MILIEU

def step_pt_milieu(y,t,h,F):
    n=np.shape(y)
    y_aux=np.zeros(n) ##y_aux represente y(n+1/2)
    z=np.zeros(n)
    y_aux=y+(h/2)*F(y,t)
    z=y+h*F(y_aux,t+h/2)
    return z

## METHODE DE HEUN

def step_heun(y,t,h,F):
    n=np.shape(y)
    z=np.zeros(n)
    y1=step_euler(y,t,h,F)
    p1=F(y1,t+h)
    p2=F(y,t)
    z=y+(1/2.)*h*(p1+p2)
    return z

## METHODE DE RUNGE-KUTTA 4

def step_rk4(y,t,h,F):
    n=np.shape(y)
    z=np.zeros(n)
    p1=F(y,t)
    y2=y+(h/2)*p1
    p2=F(y2,t+(h/2))
    y3=y+(h/2)*p2
    p3=F(y3,t+(h/2))
    y4=y+h*p3
    p4=F(y4,t+h)
    z=y+(h/6)*(p1+2*p2+2*p3+p4)
    return z


## N ETAPES

def meth_N_step(y0, t0, N, h, F, meth):
    T = [y0]
    y=y0
    t=t0
    x = t0
    for i in range(0, N):
        x = float(x+h)
        y = meth(y,t,h,F)
        T = T + [y]
    return T


def meth_N_step_theo(t0, N, h, f):
    T = [f(0, t0)]
    x = t0
    for i in range(0, N):
        x = float(x+h)
        y = f(0, x)
        T = T + [y]
    return T


### Tracer du champ des tangentes


def tangent_field(x_max,y_max,fct):
    """ Trace le champ des tangentes sur le pave [-x_max,x_max]*[-y_max,y_max] pour la fonction """
    N = 20
    X = np.linspace(-x_max, x_max, N)
    Y = np.linspace(-y_max, y_max, N)
    U = np.ones([N,N])
    V = np.ones([N,N])
    for i in range(N):
        for j in range(N):
            V[i,j] = fct([Y[i]],Y[j])
    mp.clf()
    mp.quiver(X,Y,U,V)
    mp.title('Champ de tangentes pour la fonction f(x,y) ')
    mp.xlabel('x ')
    mp.ylabel('y ')
    
    mp.savefig("champ-tangentes.png")                
    mp.show()



def test_differential(y0, t0, tf, step, f_expe, f_theo, methods_nb = 4):

    methods = [
        step_euler, "Methode d'euler",
        step_pt_milieu, "Methode du point milieu",
        step_heun, "Methode de Heun",
        step_rk4, "Methode de Runge-Kutta"]

    d = tf - t0
    N = int (d / step)

    T_theo = meth_N_step_theo(t0, N, step, f_theo)
    for i in range(0, methods_nb):
        T_expe = meth_N_step(y0, t0, N, step, f_expe, methods[2*i])
        if len(np.shape(T_theo)) == 1:
            plot_differential (T_expe, T_theo, methods[2*i+1])
        else :
            for k in range (np.shape(T_theo)[1]):
                T_expe_f = []
                T_theo_f = []
                for j in range (len(T_theo)):
                    T_expe_f = T_expe_f + [T_expe[j][k]]
                    T_theo_f = T_theo_f + [T_theo[j][k]]
                plot_differential (T_expe_f, T_theo_f, ("courbe", k+1, methods[2*i+1]))


def plot_differential (T_expe, T_theo, title=""):
    
    plt1, = mp.plot(T_expe)
    plt2, = mp.plot(T_theo)
    mp.legend([plt1, plt2], ["courbe experimentale", "courbe theorique"])
    mp.title (title)
    # mp.savefig("resolution_intro")
    mp.show()




if __name__ == "__main__":

    tangent_field(4,4,lambda x,y:np.cos(x))
    
    a = 1 << 0
    b = 1 << 1
    c = 1 << 2

    ## entrer le code du test a effectuer ##
    test = a
    methods_nb = 4
    ## ----------- ##


    if test & a:
        # test personnel
        y0 = np.array(1)
        t0 = 0
        tf = 5
        step = 0.01
        f_expe = lambda y, t: y
        f_theo = lambda y, t: exp(t)
        test_differential(y0, t0, tf, step, f_expe, f_theo, methods_nb)

    if test & b:
        # test 1
        y0 = np.array(1)
        t0 = 0
        tf = 2.5
        step = 0.001
        f_expe = lambda y, t: float(y / (1 + t**2))
        f_theo = lambda y, t: exp(np.arctan(t))
        test_differential(y0, t0, tf, step, f_expe, f_theo, methods_nb)
    
    if test & c:
        # test 2
        y0 = np.array([1, 0])
        t0 = 0
        tf = 6*np.pi
        step = 0.1
        f_expe = lambda y, t: np.array([-y[1], y[0]])
        f_theo = lambda y, t: np.array([np.cos(t), np.sin(t)])
        test_differential(y0, t0, tf, step, f_expe, f_theo, methods_nb)